/*
 * Alloy Analyzer
 * Copyright (c) 2002 Massachusetts Institute of Technology
 *
 * The Alloy Analyzer is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * The Alloy Analyzer is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with the Alloy Analyzer; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

package alloy.transform;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import alloy.ast.ASTDepthFirstVisitor;
import alloy.ast.ComprehensionExpr;
import alloy.ast.Decl;
import alloy.ast.Decls;
import alloy.ast.EmptyFormula;
import alloy.ast.EmptySetExpr;
import alloy.ast.IdentityExpr;
import alloy.ast.Leaf;
import alloy.ast.LeafId;
import alloy.ast.LiteralIntExpr;
import alloy.ast.MultiplicityExpr;
import alloy.ast.Op;
import alloy.ast.QuantifiedFormula;
import alloy.ast.SigExpr;
import alloy.ast.Signature;
import alloy.ast.SumIntExpr;
import alloy.ast.UniversalExpr;
import alloy.ast.VarCreator;
import alloy.ast.Variable;
import alloy.ast.VariableExpr;
import alloy.semantic.LocalScope;
import alloy.type.RelationType;
import alloy.type.UnionType;
import alloy.util.Dbg;

/** A visitor which sets the leaf ids for all leaves it encounters.
    This visitor is intended to be applied to the finished formula for
    a command, and does not deal with invocations, withs, and other nodes
    that are desugared before that point. */

public class SetLeafIdsVisitor extends ASTDepthFirstVisitor {

    /** current leaf */
    private LeafIdMap _curLeafIdMap;

    /** current map from leaf ids to mult exprs */
    private Map _leafIdToMult;
	
    /** The id for an {@link EmptyFormula}, if we have seen one. */
    private LeafId _emptyFormulaId;

    /** a mapping from the String representations of non-expression
        leaves to their leaf ids, for the purpose of reusing the ids.
	{@link Map} from {@link String} to {@link LeafId}. */
    private Map _nonExprLeafIdMap = new HashMap();

    /** LeafId mapping for skolem functions. */
    private static LeafIdMap _skolemLeafIdMap;
    
    /** [QualifiedName->Signature] maps sig names to their corresponding sig
     *  for all Signatures in the Spec
     **/
    private Map _sigNameToSig;
    	
    public SetLeafIdsVisitor(Map leafIdToMult_, Map sigNameToSig_) {
        _leafIdToMult = leafIdToMult_;
	_sigNameToSig = sigNameToSig_;
    }

    /** Record the LeafIdMap used by {@link SkolemizationVisitor} for this problem. */
    public void setSkolemLeafIdMap(LeafIdMap skolemLeafIdMap_) { _skolemLeafIdMap = skolemLeafIdMap_; }

    public LeafIdMap getSkolemLeafIdMap() { return _skolemLeafIdMap; }

    public static void clearState() { _skolemLeafIdMap = null; }
    
    // nodes which can create new names //

    public void visit(VarCreator creator_) {
        // save old leaf id map
        LeafIdMap oldLeafIdMap = _curLeafIdMap;
        if ((_curLeafIdMap = creator_.getLeafIdMap()) == null) {
            // generate leaf ids for newly bound names
            LocalScope localScope = creator_.getLocalScope();
            if (localScope != null) {
                _curLeafIdMap = localScope.generateLeafIdMapping(oldLeafIdMap);
            } else {
                // generate id mapping from decls
                _curLeafIdMap = _genLeafIdMapping(creator_.getDecls(), oldLeafIdMap);
            }
            creator_.setLeafIdMap(_curLeafIdMap);
        } 
	//System.out.println("Var creator " + creator_.hashCode() + ": " + creator_.nodeString());
	//System.out.println("Leaf id map: " + _curLeafIdMap);
        _updateLeafIdToMult(creator_.getDecls());
        // visit all children
        visit((alloy.ast.Node)creator_);
        // restore leaf id map
        _curLeafIdMap = oldLeafIdMap;
    }
        
    public void visit(QuantifiedFormula quantForm_) {
        visit((VarCreator)quantForm_);
    }

    public void visit(ComprehensionExpr compExpr_) {
        visit((VarCreator)compExpr_);
    }

    public void visit(SumIntExpr sumIntExpr_) {
        visit((VarCreator)sumIntExpr_);
    }

    private void _updateLeafIdToMult(Decls decls_) {
        if (_leafIdToMult != null) {
	    Iterator declIter = decls_.getDeclIter();
	    while (declIter.hasNext()) {
	        Decl decl = (Decl)declIter.next();
	        MultiplicityExpr multExpr = decl.getExpr();
	        Iterator varIter = decl.getVariables().getVariableIter();
	        while (varIter.hasNext()) {
	            LeafId id = _curLeafIdMap.lookupLeafId(((Variable)varIter.next()).nodeString());
                    if (!_leafIdToMult.containsKey(id)) {
                        _leafIdToMult.put(id, multExpr.copy());
                    }
	        }
	    }
        }
    }
    // leaf nodes //

    public void visit(SigExpr sigExpr_) {
        //Dbg.info(sigExpr_.replacednodeString());
        if (sigExpr_.getLeafId() == null) {
	    // get signature object
	    Signature sig = (Signature) _sigNameToSig.get(sigExpr_.getSig());
	    if (sig == null) {
		//System.out.println(sigExpr_.getParent().getParent().nodeString());
	    }
	    Dbg.chk(sig != null, "null signature for " + sigExpr_);
	    // get id
	    LeafId id = sig.getLeafId();
	    Dbg.chk(id != null, "null signature id for " + sigExpr_.nodeString());
	    sigExpr_.setLeafId(id);
	}
    }


    public void visit(VariableExpr variableExpr_) {
        // check to make sure leaf id hasn't already
        // been set (will have been set for skolem constants)
        // *note*: i don't know what the problem with this null check is, but it screws up 
        // many things in atomization
        //if (variableExpr_.getLeafId() == null) {
            // see if variable is bound in current mapping
            String varString = variableExpr_.getVar().getId().nodeString();
            LeafId id = (_curLeafIdMap == null) ? null : _curLeafIdMap.lookupLeafId(varString);
            if (id == null) { // not a quantified variable
                if (_skolemLeafIdMap != null) {
                    // see if it's a skolem constant
                    id = _skolemLeafIdMap.lookupLeafId(varString);
                } else {
		    //System.out.println("no skolem leaf map");
		}
                if (id == null) {
                    // variable must be a field
                    // get signature object
		    
                    Signature sig = _getSignatureObject(varString, variableExpr_.getType());
                    //Dbg.info(sig.getTypedName());
                    id = sig.getFieldLeafId(varString);
                }
            } else {
                // quantified variable, so set flag
                variableExpr_.isQuantified = true;
            }
	    if (id == null) {
		// if it's a skolem constant, its id was
		// set by the skolemization visitor
		id = variableExpr_.getLeafId();
	    }
            Dbg.chk(id != null, "no id for " + varString + " at " + variableExpr_.getLocation().toString());
            variableExpr_.setLeafId(id);
	    //}        
    }

    public void visit(LiteralIntExpr litIntExpr_) {
        _handleNonExprLeaf(litIntExpr_);
    }

    public void visit(Op op_) { 
	_handleNonExprLeaf(op_); 
    }

    public void visit(UniversalExpr univExpr_) {
        _handleNonExprLeaf(univExpr_);
    }

    public void visit(IdentityExpr idenExpr_) {
	_handleNonExprLeaf(idenExpr_);
    }

    public void visit(EmptySetExpr emptySetExpr_) {
	_handleNonExprLeaf(emptySetExpr_);
    }




    public void visit(EmptyFormula emptyFormula_) {
	if (_emptyFormulaId == null)
	    _emptyFormulaId = new LeafId();
	emptyFormula_.setLeafId(_emptyFormulaId);
    }

    private void _handleNonExprLeaf(Leaf leaf_) {
        if (leaf_.getLeafId() == null) {
            LeafId id = null;
            // see if id has already been created
            String leafKey = leaf_.getClass().toString() + "##" + leaf_.nodeString();
            if ((id = (LeafId)_nonExprLeafIdMap.get(leafKey)) == null) {
                // create new id if necessary
                _nonExprLeafIdMap.put(leafKey, id = new LeafId());
            }
            // set leaf id
            leaf_.setLeafId(id);
        }
    }

    private static boolean _sigHasField(Signature sig, String field) {
	Iterator decls = (sig).getDecls().getDeclIter();
	Decl decl;

	while(decls.hasNext()){
	    decl = (Decl)decls.next();
	    if (!decl.isEmpty()) {
		Iterator varIter = decl.getVariables().getVariableIter();
		while (varIter.hasNext()) {
		    Variable curVar = (Variable)varIter.next();
		    String varName = curVar.getId().nodeString();
		    if (varName.equals(field)) {
			return true;
		    }
		}
	    }
	}
	
	return false;

    }

    private Collection _getSigsForField(String fieldName_) {
	Iterator sigIter = _sigNameToSig.values().iterator();
	List sigs = new ArrayList();
	while (sigIter.hasNext()) {
	    Signature sig = (Signature) sigIter.next();
	    if (_sigHasField(sig, fieldName_)) 
		sigs.add(sig);
	}

	return sigs;
    }

    /**
     * @return a Signature object which has a field named fieldName_ of type fieldType_.
     * If no such signature exists, returns null.  Note that fieldName_ and fieldType_
     * uniquely determine a signature since no two fields with intersecting first column
     * are allowed to have the same name.
     **/
    private Signature _getSignatureObject(String fieldName_, UnionType fieldType_) {

	// get all Signatures that have a field named fieldName_
	
	Collection sigs = _getSigsForField(fieldName_);
	
	//System.out.println("fieldName_ " + fieldName_ + " : " + fieldType_);
	Dbg.check(!sigs.isEmpty(),"No sig with field named " + fieldName_ + " exists.");
	
	// construct a UnionType that is a sum of first BasicTypes in each
	// RelationType contained in fieldType_
	Iterator iter = fieldType_.getTypeIter();
	UnionType firstColType = UnionType.empty(1);
	while (iter.hasNext()) {
	    firstColType = firstColType.union(new UnionType(new RelationType(((RelationType) iter.next()).getBasicTypeAt(0))));
	}
	Dbg.chk(!firstColType.isEmpty(),"should not reach here");
	
	// iterate through sigs and return the one whose type intersects the
	// constructed firstColType (there can be only one such Signature)
	iter = sigs.iterator();
	while (iter.hasNext()) {
	    Signature curSig = (Signature) iter.next();
	    if (!firstColType.intersect(curSig.getUnionType()).isEmpty())
		return curSig;
	    
	}
	
	throw new RuntimeException("No sig with field " + fieldName_ +
				   " has type that intersects " + fieldType_);


    }

    
    private LeafIdMap _genLeafIdMapping(Decls decls_, LeafIdMap parent_) {
        LeafIdMap map = new LeafIdMap(parent_);
	Iterator declIter = decls_.getDeclIter();
	while(declIter.hasNext()) {
	    Decl decl = (Decl)declIter.next();
	    Iterator vars = decl.getVariables().getVariableIter();
	    while(vars.hasNext()) {
		Variable var = (Variable)vars.next();
                //System.out.println("generating id for " + var.orignodeString());
                map.put(var.nodeString(), new LeafId());
            }
        }
        return map;
    }

    public String toString() { return _skolemLeafIdMap==null ? "[no map]" : _skolemLeafIdMap.toString(); }
    
    public Map getNonExprLeafIdMap() {
        return _nonExprLeafIdMap;
    }
    
    public LeafId getEmptyFormulaLeafId() {
        return _emptyFormulaId;
    }
}




